from logging import warn
from tensorboard.backend.event_processing import event_accumulator
import matplotlib.pyplot as plt
import os

from torch.utils import data

# ea0 = event_accumulator.EventAccumulator("semi_logits_result_fixmatch_sfclr100.0/recognition36/resnet_MyResSA(8_64)_pos-embed/5shot-5way_CosineLR_unlabel30_lambdau-0.0_mse-10_bs-4_mu-4_threshold-0.8_remove-interleave/tf/events.out.tfevents.1639622342.yuan-System-Product-Name.28736.0", size_guidance={event_accumulator.SCALARS:180000})
# ea0_1 = event_accumulator.EventAccumulator("semi_logits_result_fixmatch_sfclr100.0/recognition36/resnet_MyResSA(8_64)_pos-embed/5shot-5way_CosineLR_unlabel30_lambdau-0.1_mse-10_bs-4_mu-4_threshold-0.8_remove-interleave/tf/events.out.tfevents.1640679262.yuan-System-Product-Name.17538.0", size_guidance={event_accumulator.SCALARS:180000})
ea0_2 = event_accumulator.EventAccumulator("semi_logits_result_rotation_sfclr100.0/recognition36/resnet_MyResSA(8_64)_pos-embed_relu/5shot-5way_CosineLR_unlabel30_rotw-0.5_mse-10_rotbs-4_rotmu-1_rotlr-0.1/tf/events.out.tfevents.1640902646.yuan-System-Product-Name.27105.0", size_guidance={event_accumulator.SCALARS:180000})
eas = [
    # ea0, 
    # ea0_1, 
    ea0_2
    ]

for ea in eas:
    ea.Reload()
    keys = ea.scalars.Keys()

    unlabel_loss = ea.scalars.Items(keys[0])
    w = ea.scalars.Items(keys[1])
    loss = ea.scalars.Items(keys[2])

    unlabel_loss = [(i.step, i.value) for i in unlabel_loss]
    w = [(i.step, i.value) for i in w]
    loss = [(i.step, i.value) for i in loss]

    unlabel_loss_0, unlabel_loss_1 = zip(*unlabel_loss)
    w_0, w_1 = zip(*w)
    loss_0, loss_1 = zip(*loss)

    index = 30
    # lenth = len(unlabel_loss_1)
    counter = range(index)
    # assert len(unlabel_loss_1) == len(counter)

    plt.figure()
    plt.subplot(3,1,1)
    plt.plot(counter, unlabel_loss_1[-index:], label=keys[0])
    plt.legend()
    plt.subplot(3,1,2)
    plt.plot(counter, w_1[-index:], label=keys[1])
    plt.legend()
    plt.subplot(3,1,3)
    plt.plot(counter, loss_1[-index:], label=keys[2])
    plt.legend()

plt.show()